package com.mycompany.myplugin;

import com.intellij.openapi.actionSystem.AnAction;
import com.intellij.openapi.actionSystem.AnActionEvent;
import com.intellij.openapi.editor.Editor;
import com.intellij.openapi.project.Project;

import com.intellij.psi.*;
import com.intellij.psi.util.*;

import okhttp3.*;
import org.json.JSONArray;
import org.json.JSONException;
import org.json.JSONObject;

import java.io.IOException;
import java.util.concurrent.CompletableFuture;

public class MyPluginAction extends AnAction {
    @Override
    public void actionPerformed(AnActionEvent e) {
        Project project = e.getProject();
        if (project != null) {
            Editor editor = (Editor) e.getDataContext().getData("editor");
            if (editor != null) {
                PsiMethod psiMethod = getSelectedMethod(editor);
                if (psiMethod != null) {
                    String methodName = psiMethod.getName();

                    // 调用ChatGPT API来生成注释
                    CompletableFuture<String> commentFuture = CompletableFuture.supplyAsync(
                            () -> generateComment(methodName)
                    );

                    // 在生成注释之前，你可能需要显示一个加载提示或进度条
                    commentFuture.thenAccept(comment -> {
                        // 将生成的注释插入到方法体内
                        insertComment(psiMethod, comment);
                        // 在插入注释后，你可能需要关闭加载提示或进度条
                    });
                }
            }
        }
    }

    private PsiMethod getSelectedMethod(Editor editor) {
        // 获取当前光标所在位置的元素
        PsiFile psiFile = PsiDocumentManager.getInstance(editor.getProject()).getPsiFile(editor.getDocument());
        PsiElement element = psiFile.findElementAt(editor.getCaretModel().getOffset());

        // 利用 PsiTreeUtil 获取当前光标所在位置的方法
        PsiMethod method = PsiTreeUtil.getParentOfType(element, PsiMethod.class);

        return method;
    }

    private String generateComment(String methodName) {
        // 调用 ChatGPT API 来生成注释
        String endpoint = "https://api.openai.com/v1/engines/davinci-codex/completions";
        String apiKey = "API_KEY";

        OkHttpClient client = new OkHttpClient();

        MediaType mediaType = MediaType.parse("application/json");
        RequestBody body = RequestBody.create(mediaType, "{\"prompt\": \"Write a comment for method " + methodName + ".\", \"max_tokens\": 100}");
        Request request = new Request.Builder()
                .url(endpoint)
                .post(body)
                .addHeader("Content-Type", "application/json")
                .addHeader("Authorization", "Bearer " + apiKey)
                .build();

        try {
            Response response = client.newCall(request).execute();

            if (response.isSuccessful()) {
                String responseStr = response.body().string();
                // 解析 API 响应并提取注释
                String comment = extractCommentFromResponse(responseStr);
                return comment;
            } else {
                System.out.println("API request unsuccessful. Response code: " + response.code() + ", message: " + response.message());
            }
        } catch (IOException e) {
            e.printStackTrace();
        }

        return "";
    }

    private String extractCommentFromResponse(String responseStr) {
        // 解析 API 响应并提取生成的注释
        // 根据 ChatGPT API 的响应格式进行解析

        // 解析示例：
        // 假设 API 响应是 JSON 格式，包含一个名为 "choices" 的数组，其中每个元素包含生成的文本。
        // 简单示例，仅提取第一个元素作为注释。

        try {
            JSONObject jsonObject = new JSONObject(responseStr);
            JSONArray choicesArray = jsonObject.getJSONArray("choices");

            if (choicesArray.length() > 0) {
                JSONObject choiceObject = choicesArray.getJSONObject(0);
                String comment = choiceObject.getString("text");

                return comment;
            }
        } catch (JSONException e) {
            e.printStackTrace();
        }

        return null;
    }

    private void insertComment(PsiMethod psiMethod, String comment) {
        Project project = psiMethod.getProject();
        PsiElementFactory elementFactory = JavaPsiFacade.getElementFactory(project);

        PsiStatement[] statements = psiMethod.getBody().getStatements();

        if (statements.length > 0) {
            PsiComment psiComment = elementFactory.createCommentFromText("//" + comment, null);
            psiMethod.getBody().addBefore(psiComment, statements[0]);
        } else {
            PsiComment psiComment = elementFactory.createCommentFromText("//" + comment, null);
            psiMethod.getBody().add(psiComment);
        }
    }
}